#include "nv_utils.h"
#include "logger.h"
using namespace std;

_TRT_INFER_BEGIN
namespace CUDATools
{

    bool checkCudaRuntime(cudaError_t e, const char *call, int line, const char *file)
    {
        if (e != cudaSuccess)
        {
            LOGE << "CUDA Runtime error: " << call << " # " << cudaGetErrorString(e)
                 << " # " << cudaGetErrorName(e) << " in file " << file << " at line " << line;
            return false;
        }
        return true;
    }

    bool check_device_id(int device_id)
    {
        int device_count = -1;
        checkRuntime(cudaGetDeviceCount(&device_count));
        if (device_id < 0 || device_id >= device_count)
        {
            LOGE << "Invalid device id: " << device_id << ", count = " << device_count;
            return false;
        }
        return true;
    }

    string device_description()
    {

        cudaDeviceProp prop;
        size_t free_mem, total_mem;
        int device_id = 0;

        checkRuntime(cudaGetDevice(&device_id));
        checkRuntime(cudaGetDeviceProperties(&prop, device_id));
        checkRuntime(cudaMemGetInfo(&free_mem, &total_mem));

        return format(
            "[ID %d]<%s>[arch %d.%d][GMEM %.2f GB/%.2f GB]",
            device_id, prop.name, prop.major, prop.minor,
            free_mem / 1024.0f / 1024.0f / 1024.0f,
            total_mem / 1024.0f / 1024.0f / 1024.0f);
    }

    void set_device_id(int device_id)
    {
        if (device_id == -1)
            return;
        checkRuntime(cudaSetDevice(device_id));
    }

    int current_device_id()
    {
        int device_id = 0;
        checkRuntime(cudaGetDevice(&device_id));
        return device_id;
    }

    dim3 grid_dims(int numJobs)
    {
        int numBlockThreads = numJobs < GPU_BLOCK_THREADS ? numJobs : GPU_BLOCK_THREADS;
        return dim3((numJobs + numBlockThreads - 1) / (float)numBlockThreads);
    }

    dim3 block_dims(int numJobs)
    {
        return numJobs < GPU_BLOCK_THREADS ? numJobs : GPU_BLOCK_THREADS;
    }

    AutoDevice::AutoDevice(int device_id)
    {
        cudaGetDevice(&old_);
        checkRuntime(cudaSetDevice(device_id));
    }

    AutoDevice::~AutoDevice()
    {
        checkRuntime(cudaSetDevice(old_));
    }

} // namespace CUDATools

inline static int check_and_trans_device_id(int device_id)
{
    if (device_id != CURRENT_DEVICE_ID)
    {
        CUDATools::check_device_id(device_id);
        return device_id;
    }

    checkRuntime(cudaGetDevice(&device_id));
    return device_id;
}

_TRT_INFER_END